package tr.helper;

import com.jcraft.jsch.JSch;
import com.jcraft.jsch.JSchException;
import com.jcraft.jsch.Session;
import java.util.HashMap;
import java.util.Map;
import org.apache.commons.pool2.ObjectPool;
import org.apache.commons.pool2.PooledObject;
import org.apache.commons.pool2.impl.DefaultPooledObject;
import org.apache.commons.pool2.impl.GenericObjectPool;

/**
 * tr-service
 *
 * @author MrYe
 * @date 2024-07-29 09:30
 */
public class SessionPoolManager {

    private Map<String, ObjectPool<Session>> sessionPools;

    public SessionPoolManager() {
        sessionPools = new HashMap<>();
    }

    public Session getSession(String serverKey, String host, int port, String username, String password) {
        ObjectPool<Session> pool = sessionPools.computeIfAbsent(serverKey, key -> createSessionPool(host, port, username, password));
        // 从特定服务器的连接池中获取Session
        try {
            return pool.borrowObject();
        } catch (Exception e) {
            e.printStackTrace();
            return null;
        }
    }

    public void releaseSession(String serverKey, Session session) {
        ObjectPool<Session> pool = sessionPools.get(serverKey);
        if (pool != null) {
            try {
                // 将Session放回连接池
                pool.returnObject(session);
            } catch (Exception e) {
                e.printStackTrace();
            }
        }
    }

    public void closeAllSessions() {
        for (ObjectPool<Session> pool : sessionPools.values()) {
            try {
                pool.close(); // 关闭所有连接池
            } catch (Exception e) {
                e.printStackTrace();
            }
        }
        sessionPools.clear();
    }

    private ObjectPool<Session> createSessionPool(String host, int port, String username, String password) {
        return new GenericObjectPool<>(new SessionFactory(host, port, username, password));
    }

    // 自定义SessionFactory实现
    private static class SessionFactory extends org.apache.commons.pool2.BasePooledObjectFactory<Session> {

        private String host;
        private int port;
        private String username;
        private String password;

        public SessionFactory(String host, int port, String username, String password) {
            this.host = host;
            this.port = port;
            this.username = username;
            this.password = password;
        }

        @Override
        public Session create() throws JSchException {
            JSch jSch = new JSch();
            Session session = jSch.getSession(username, host, port);
            session.setPassword(password);
            session.setConfig("StrictHostKeyChecking", "no");
            session.connect();
            return session;
        }

        @Override
        public void destroyObject(org.apache.commons.pool2.PooledObject<Session> p) {
            Session session = p.getObject();
            session.disconnect();
        }

        @Override
        public boolean validateObject(org.apache.commons.pool2.PooledObject<Session> p) {
            Session session = p.getObject();
            return session.isConnected();
        }

        @Override
        public PooledObject<Session> wrap(Session session) {
            return new DefaultPooledObject<>(session);
        }

    }

}
